library(dplyr)
library(ggplot2)
library(rstan)

load("fit.Rdata")
mdlAut <- mdl
polexAut <- polex
fitAut <- fit

load("fit.Rdata")
mdlDem <- mdl
polexDem <- polex
fitDem <- fit

load("fit.Rdata")

ests <- function(fit) {
  s <- summary(fit)$summary
  s[-grep("gamma|lp", rownames(s)), c("mean", "sd")]
}

ests(fit)
ests(fitDem)
ests(fitAut)

p <- merge(polex, polexAut %>% select(ruler, starty, gamma), by = c("ruler", "starty"), all.x = TRUE, all.y = TRUE, suffixes = c("", ".aut"))

p <- merge(p, polexDem %>% select(ruler, starty, gamma), by = c("ruler", "starty"), all.x = TRUE, all.y = TRUE, suffixes = c("", ".dem"))

p$gamma.autdem <- ifelse(is.na(p$gamma.aut), p$gamma.dem, p$gamma.aut)

head(p)

pdf("compare_separate_vs_joint_model.pdf", width = 11, height = 7)

p %>%
  filter(!is.na(polity2)) %>%
  ggplot(aes(x = gamma, y = gamma.autdem,
             col = ifelse(polity2 < 6, "Autocracy", "Democracy"))) +
  geom_point(alpha = .5) +
  theme_minimal() +
  labs(x = "Gamma (full model)",
       y = "Gamma (separate models)",
       col = "Polity IV") +
  geom_abline(intercept = 0, slope = 1, col = "black", lty = "dashed") +
  theme(legend.position = "bottom") +
  scale_color_manual(values = c("black", "darkgray"))

p %>%
  filter(!is.na(polity2) & year >= 1950) %>%
  ggplot(aes(x = year, y = gamma,
             col = ifelse(polity2 < 6, "Autocracy", "Democracy"))) +
  geom_point(alpha = .3) +
  geom_smooth(se = FALSE) +
  theme_minimal() +
  labs(x = "Year",
       y = "Gamma (full model)",
       col = "Polity IV") +
  theme(legend.position = "bottom") +
  scale_color_manual(values = c("darkgray", "black"))

p %>%
  filter(!is.na(polity2) & year >= 1950) %>%
  ggplot(aes(x = year, y = gamma.autdem,
             col = ifelse(polity2 < 6, "Autocracy", "Democracy"))) +
  geom_point(alpha = .3) +
  geom_smooth(se = FALSE) +
  theme_minimal() +
  labs(x = "Year",
       y = "Gamma (separate models)",
       col = "Polity IV") +
  theme(legend.position = "bottom") +
  scale_color_manual(values = c("darkgray", "black"))

dev.off()
